
from traits.api               import HasTraits, Range, Array, Instance, Enum, on_trait_change, Bool, Button #SingletonHasTraits, Trait, Property, String, Float, Int, cached_property, Code, List, NO_COMPARE
from traitsui.api            import View, Item, HGroup #, Group, HGroup, Tabbed, EnumEditor
from enable.api               import ComponentEditor #, Component
from chaco.api                import ArrayPlotData, Plot #, ScatterPlot, CMapImagePlot, 
from traitsui.file_dialog    import save_file
from traitsui.menu           import Action, Menu, MenuBar

from tools.utility                      import CloseHandler, ThreadedMeasurementMixin, GetSetItemsMixin

from TimeTagger                         import Pulsed, Countrate


from datetime                           import datetime
from time                               import sleep
from threading                          import currentThread

import numpy  

class Autocorrelation(HasTraits, ThreadedMeasurementMixin, GetSetItemsMixin):
    
    WindowSize = Range(low=10, high=10000, value=500, desc='WindowSize in the time domain  [ns]', label='WindowSize')
    chan1 = Enum(0,1,2,3, desc="the trigger channel", label="Channel 1")
    chan2 = Enum(0,1,2,3, desc="the signal channel", label="Channel 2")
    remove_zero = Bool(desc="remove the very high value at position zero", label="remove zero")
    restart = Button(desc="restart the measuremen", label="restart")

    C = Array()
    T = Array()
    TracePlot = Instance( Plot, transient=True )
    TraceData = Instance( ArrayPlotData, transient=True )
    
    get_set_items=['WindowSize', 'chan1', 'chan2', 'C', 'T', 'remove_zero']
    
    def __init__(self):
        ThreadedMeasurementMixin.__init__(self)
        HasTraits.__init__(self)
        
        self.bin_with = 1.25
        self._update_pulsed()
        
        self.start()
        
    @on_trait_change('chan1,chan2,WindowSize,restart')    
    def _update_pulsed(self):
        self.p1 = Pulsed(int(self.WindowSize/(2*self.bin_with)),1,1,self.chan1,self.chan2)
        self.p2 = Pulsed(int(self.WindowSize/(2*self.bin_with)),1,1,self.chan2,self.chan1) 
        self.count_rate_1 = Countrate(self.chan1)
        self.count_rate_2 = Countrate(self.chan2)
        self.started = datetime.now()
    
    def run(self):
        while True:
            data1 = self.p1.getData()
            data2 = self.p2.getData()
            if self.remove_zero:
                data = numpy.append(numpy.append(data1[0][-1:0:-1], 0),data2[0][1:])
            else:
                data = numpy.append(numpy.append(data1[0][-1:0:-1], max(data1[0][0],data2[0][0])),data2[0][1:])
        
            timediff = datetime.now() - self.started
            timediff = (timediff.seconds + timediff.days * 24 * 3600) 
            
            factor = self.count_rate_1.getData()*self.count_rate_2.getData()*self.bin_with*10**-9*timediff
            if factor > 0:
                self.C = data / (self.count_rate_1.getData()*self.count_rate_2.getData()*self.bin_with*10**-9*timediff)
            else:
                self.C = self._C_default()
        
            if currentThread().stop_request.isSet():
                break
            sleep(1)

    def savePlot(self,filename=None):
        if filename is None:
            filename = pi3d.get_filename()+"_autocorrelation_plot.png"
        else:
            filename = pi3d.get_filename(filename)
        print filename
        self.save_figure(self.TracePlot, filename)
    
    def _remove_zero_default(self):
        return False
    
    def _C_default(self):
        return numpy.zeros((int(self.WindowSize/(2*self.bin_with))*2-1,))
    
    def _T_default(self):
        return self.bin_with*numpy.arange(-int(self.WindowSize/(2*self.bin_with))+1,int(self.WindowSize/(2*self.bin_with)))
    
    def _T_changed(self):
        self.TraceData.set_data('t', self.T)

    def _C_changed(self):
        self.TraceData.set_data('y', self.C)
        self.TracePlot.request_redraw()
    
    def _chan1_default(self):
        return 0
    
    def _chan2_default(self):
        return 1 
        
    def _WindowSize_changed(self):
        self.C = self._C_default()
        self.T = self._T_default()
        
    def _TraceData_default(self):
        return ArrayPlotData(t=self.T, y=self.C)
    
    def _TracePlot_default(self):
        plot = Plot(self.TraceData, width=500, height=500, resizable='hv')
        plot.plot(('t','y'), type='line', color='blue')        
        return plot
    
    view = View( 
        Item('TracePlot', editor=ComponentEditor(), show_label=False),
        HGroup(
            Item('WindowSize'),
            Item('chan1'),
            Item('chan2'),
            Item('remove_zero'),
            Item('restart', show_label=False)),
        menubar = MenuBar(
            Menu(
                Action(action='savePlot', name='SavePlot (.png)'),
                Action(action='save', name='Save (.pyd or .pys)'),
                Action(action='load', name='Load'),
                Action(action='_on_close', name='Quit'),
            name='File')),
        title='Autocorrelation', width=900, height=800, buttons=[], resizable=True,
        handler=CloseHandler)

